import jieba
import torch
import re
import os
import numpy as np
from tqdm import tqdm
import pandas as pd
from paddlenlp import Taskflow
from transformers import XLNetTokenizer, XLNetModel, BertTokenizer, BertModel
from torch.utils.data import DataLoader, TensorDataset
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# 数据预处理并分词
def pre_text(text):
    text = text.replace('！','').replace('，','').replace('。','').replace('”','').replace('“','').replace('-','').replace('？','').replace('：','')  # 将标点符号处理掉
    return jieba.lcut(text)

def pre_text_2(text):
    text = text.replace('！','').replace('，','').replace('。','').replace('”','').replace('“','').replace('-','').replace('？','').replace('：','') # 将标点符号处理掉
    cleaned_text = re.sub(r'\s+', ' ', text)
    return cleaned_text.strip()

# 清理“病情描述”中的无关数字和多余空格
def clean_description(description):
    # 使用正则表达式去掉数字和前后的空格
    cleaned = re.sub(r'^\d+\s*', '', description)  # 去掉开头的数字和空格
    return cleaned.strip()  # 去除两端的多余空格

def get_medical_result(text):
        # 获取实体
        ner = Taskflow("ner", model="medical")
        results = ner(text)

        # 进行实体筛选ner任务并富集关键词

        final_word = []
        for result in results:
            if result[1] in ['疾病损伤类', '术语类_生物体', '个性特征', '物体类_概念', '场景事件']:
                final_word.append(result[0])
   
        return final_word

def hospitalgpt_reply(text, tokenizer, model_xlnet, model):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    standard_departments = [
        "呼吸内科", "心血管内科", "消化内科", "内分泌科", "神经内科", 
        "普通外科", "心胸外科", "神经外科", "骨科", "泌尿外科", 
        "烧伤整形科", "妇科", "产科", "普通儿科", "新生儿科", 
        "眼科", "耳鼻喉科", "皮肤科", "精神心理科", "肿瘤科", 
        "感染科", "康复科", "中医科", "放射科", "急诊"
    ]

    inputs = tokenizer(
        text, padding='max_length', truncation=True, return_tensors='pt', max_length=512
    ).to(device)

    text_vectors = []
    with torch.no_grad():
        # 使用BERT模型进行嵌入
        outputs = model_xlnet(**inputs)
        # 获取最后一个隐藏层表示,其形状为 [batch_size, sequence_length, hidden_size]
        last_hidden_state = outputs.last_hidden_state
        #print(last_hidden_state.shape)
        token_embeddings = torch.squeeze(last_hidden_state, dim=0)
        #token_embeddings = last_hidden_state

        text_vectors.append(token_embeddings.cpu().detach().numpy())
    
    text_vectors = np.array(text_vectors)
    text_vectors = torch.tensor(text_vectors, dtype=torch.float32).to(device)

    with torch.no_grad():
        outputs_model = model(text_vectors)

        if not any(word in text for word in ['孩', '婴']):
            outputs_model[:, 13:15] = 0.0
            #if any(word in text for word in ['发烧']):

        _, predicted = torch.max(outputs_model, dim=1)
        predicted = predicted.item()
        department_out = standard_departments[predicted]
    
    # 获取关键词
    medical_result = get_medical_result(text)

    # 获得回复
    key_word = '，'.join(medical_result[:])
    medical_result_dialog = '提取到您描述中的关键词：'+key_word

    # 简单地情感分析
    word_bad = ['非常', '死', '救', '命', '很']
    flag_call_120 = False

    if any(word in text for word in word_bad):
        reply = '感受到您的症状很严重，为您推荐：'+ department_out + '，如果您实在感觉很难受，请前往急诊或者拨打120。'
        flag_call_120 = True
    else:
        reply_word = ['好的，已读取您的症状描述，为您推荐：', '滴滴，为您推荐以下科室：', '已收到，为您推荐以下科室：']
        num = np.random.choice(len(reply_word), 1, replace=False)[0]
        #print(num)
        reply = reply_word[num]
        reply = reply + department_out + '，如果您有其他症状，请继续描述。'
    
    return medical_result_dialog, reply, medical_result, department_out, flag_call_120
